from langchain_core.prompts import ChatPromptTemplate  # 构建prompt
from langchain_ollama import ChatOllama  # 使用ollama进行回答
from langchain_core.output_parsers import StrOutputParser  # 输出格式化
import QueryVectorDB


def chat(question):
    dbResult = QueryVectorDB.queryVectorDB(question)

    # prompt模板
    prompt = ChatPromptTemplate.from_template(""" 
    请根据已知信息回答问题
    
    ## 已知信息：
    {context}
    
    ## 问题：
    {question}
    
    ## 限制:
    - 回答时请注意问题和已知信息的相关性。
    - 只回答与已知信息相关的问题，拒绝回答无关话题。
    - 所输出的内容必须逻辑清晰、有条理。
    - 回答内容需基于已知信息，若已知信息中无相关内容，需明确告知用户无对应答案。
    """)

    llm = ChatOllama(model="deepseek-r1:1.5b")  # 绑定大模型

    output_parser = StrOutputParser()  # 过滤不需要的输出信息

    rag_chain = prompt | llm | output_parser  # 绑定之前的config

    result = rag_chain.invoke({"question": question, "context": dbResult})
    return result


result = chat("华为公司员工的工作时间是几点到几点？")
print("deepseek回答1：\n", result)
result = chat("在华为工作了13年的员工能享受几天带薪年假？")
print("deepseek回答2：\n", result)
result = chat("华为公司有哪些保密规定？")
print("deepseek回答3：\n", result)
result = chat("华为Mate系统系列手机的特点是什么？")
print("deepseek回答4：\n", result)
result = chat("如何使用ollama离线部署deepseek？")
print("deepseek回答5：\n", result)
